import os
import random
import torch
import numpy as np
from torch.utils.data import DataLoader

# 导入自定义模块
from models.unet import UNet
from models.diffusion import DiffusionModel  # 使用新的RePaint模型
from data.InpaintingDataset import InpaintingDataset
from trainers.ddpm_trainer import DDPMTrainer
from utils.checkpoint import load_checkpoint
from utils.logging_utils import setup_logging, log_args
from config import get_args

def main():
    """主函数，处理参数并开始训练"""
    # 获取参数
    args = get_args()
    
    # 设置日志
    logger = setup_logging()
    log_args(args)
    
    # 设置随机种子以便复现
    if args.seed is not None:
        torch.manual_seed(args.seed)
        np.random.seed(args.seed)
        random.seed(args.seed)
    
    # 设置设备
    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    logger.info(f"Using device: {device}")
    
    # 创建数据集和数据加载器
    train_dataset = InpaintingDataset(
        root_dir=args.data_dir,
        mode='train',
        image_size=args.image_size,
        data_percentage=args.train_percentage
    )
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=True,
        drop_last=True
    )
    
    # 创建验证数据集(如果需要)
    eval_loader = None
    if args.do_eval:
        eval_dataset = InpaintingDataset(
            root_dir=args.data_dir,
            mode='val',
            image_size=args.image_size,
            data_percentage=args.val_percentage
        )
        
        eval_loader = DataLoader(
            eval_dataset,
            batch_size=args.eval_batch_size,
            shuffle=False,
            num_workers=args.num_workers,
            pin_memory=True
        )
    
    # 初始化模型
    unet = UNet(
        c_in=3,
        c_out=3,
        time_dim=args.time_embedding_dim,
        device=device
    ).to(device)
    
    # 使用RePaint扩散模型
    diffusion = DiffusionModel(
        model=unet,
        beta_start=args.beta_start,
        beta_end=args.beta_end,
        timesteps=args.diffusion_steps,
        device=device
    ).to(device)
    
    # 初始化优化器
    optimizer = torch.optim.AdamW(
        diffusion.parameters(),
        lr=args.learning_rate,
        weight_decay=args.weight_decay
    )
    
    # 学习率调度器
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=args.num_epochs * len(train_loader),
        eta_min=args.min_learning_rate
    )
    
    # 创建训练器
    trainer = DDPMTrainer(
        model=diffusion,
        train_loader=train_loader,
        eval_loader=eval_loader,
        optimizer=optimizer,
        scheduler=scheduler,
        device=device,
        args=args
    )
    
    # 加载检查点（如果有）
    if args.resume_from_checkpoint and os.path.exists(args.resume_from_checkpoint):
        trainer.start_epoch, trainer.global_step = load_checkpoint(
            diffusion, 
            optimizer, 
            args.resume_from_checkpoint
        )
    
    # 开始训练
    trainer.train()

if __name__ == "__main__":
    main()